{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ac9ca6c1-9fec-4904-8066-1ef66d67e9af",
   "metadata": {},
   "source": [
    "# Visualize NWM\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98567175-d71b-4b45-822f-59b3e44bbcdd",
   "metadata": {},
   "source": [
    "## Load model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a934873-7c87-43f8-a690-218c1fd76b08",
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "import torch\n",
    "torch.backends.cuda.matmul.allow_tf32 = True\n",
    "torch.backends.cudnn.allow_tf32 = True\n",
    "from PIL import Image\n",
    "import requests\n",
    "from io import BytesIO\n",
    "import numpy as np\n",
    "from IPython.display import display, HTML\n",
    "import ipywidgets as widgets\n",
    "from diffusers.models import AutoencoderKL\n",
    "\n",
    "from diffusion import create_diffusion\n",
    "from isolated_nwm_infer import model_forward_wrapper\n",
    "from misc import transform\n",
    "from models import CDiT_models\n",
    "from datasets import TrainingDataset\n",
    "\n",
    "EXP_NAME = 'nwm_cdit_xl'\n",
    "MODEL_PATH = f'logs/{EXP_NAME}/checkpoints/0100000.pth.tar'\n",
    "\n",
    "with open(\"config/data_config.yaml\", \"r\") as f:\n",
    "    default_config = yaml.safe_load(f)\n",
    "config = default_config\n",
    "\n",
    "with open(f'config/{EXP_NAME}.yaml', \"r\") as f:\n",
    "    user_config = yaml.safe_load(f)\n",
    "config.update(user_config)\n",
    "latent_size = config['image_size'] // 8\n",
    "\n",
    "print(\"loading model\")\n",
    "model = CDiT_models[config['model']](input_size=latent_size, context_size=config['context_size'])\n",
    "ckp = torch.load(MODEL_PATH, map_location='cpu', weights_only=False) \n",
    "\n",
    "print(model.load_state_dict(ckp[\"ema\"], strict=True))\n",
    "model.eval()\n",
    "device = 'cuda:0'\n",
    "model.to(device)\n",
    "model = torch.compile(model)\n",
    "\n",
    "diffusion = create_diffusion(str(250))\n",
    "vae = AutoencoderKL.from_pretrained(f\"stabilityai/sd-vae-ft-ema\").to(device)\n",
    "latent_size = config['image_size'] // 8\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b496bed-ae5e-4478-b9ac-1e59aa1e0a98",
   "metadata": {},
   "source": [
    "## Choose starting image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "547fbb4c-050c-4ad5-b4ef-90649d35dfd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def url_to_pil_image(url):\n",
    "    response = requests.get(url)\n",
    "    img = Image.open(BytesIO(response.content)).convert('RGB')\n",
    "    return img\n",
    "\n",
    "def load_internet_image(url):\n",
    "    from torchvision import transforms\n",
    "    _transform = transforms.Compose([\n",
    "        transforms.Resize((224, 224)),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),\n",
    "    ])\n",
    "    img = url_to_pil_image(url)\n",
    "    x_start = _transform(img)\n",
    "    return x_start.unsqueeze(0).expand(config['context_size'], x_start.shape[0], x_start.shape[1], x_start.shape[2])\n",
    "    \n",
    "# Jupyter Notebook Cell\n",
    "\n",
    "\n",
    "# List of image links\n",
    "# For unknown environments, please see https://github.com/facebookresearch/nwm/issues/7\n",
    "image_links = [\n",
    "    'https://raw.githubusercontent.com/amirbar/amirbar.github.io/refs/heads/master/images/recon.png',\n",
    "    'https://raw.githubusercontent.com/amirbar/amirbar.github.io/refs/heads/master/images/scand.png',\n",
    "    'https://raw.githubusercontent.com/amirbar/amirbar.github.io/refs/heads/master/images/sacson.png',\n",
    "    'https://raw.githubusercontent.com/amirbar/amirbar.github.io/refs/heads/master/images/tartan.png'\n",
    "]\n",
    "\n",
    "# Output widget to hold the selected link\n",
    "output = widgets.Output()\n",
    "x_start_link = None  # This will hold the selected link\n",
    "\n",
    "# Function to handle image click\n",
    "def on_image_click(link):\n",
    "    global x_start_link\n",
    "    x_start_link = link\n",
    "    with output:\n",
    "        output.clear_output()\n",
    "        print(f\"Selected image link:\\n{x_start_link}\")\n",
    "\n",
    "# Create HBox of images\n",
    "image_buttons = []\n",
    "for link in image_links:\n",
    "    img = widgets.Button(\n",
    "        description='click',\n",
    "        layout=widgets.Layout(width='150px', height='20px', padding='0'),\n",
    "        style={'button_color': 'lightgray'}\n",
    "    )\n",
    "\n",
    "    img._dom_classes += ('image-button',)\n",
    "    img_link = link  # capture current link in closure\n",
    "\n",
    "    def on_click(b, link=img_link):\n",
    "        on_image_click(link)\n",
    "\n",
    "    img.on_click(on_click)\n",
    "\n",
    "    # Embed image using HTML style\n",
    "    img_html = f'<img src=\"{link}\" width=\"150px\" height=\"150px\">'\n",
    "    img_html_widget = widgets.HTML(value=img_html)\n",
    "    image_buttons.append(widgets.VBox([img_html_widget, img]))\n",
    "\n",
    "# Display the gallery\n",
    "display(widgets.HBox(image_buttons))\n",
    "display(output)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acbb2fb1-47ba-480f-8d23-15411def6c2e",
   "metadata": {},
   "source": [
    "## Visualize navigation commands"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84e26069-8f2b-4ac1-a96d-2778bfa1a89a",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_start = load_internet_image(x_start_link)\n",
    "\n",
    "commands = {\n",
    "    'Forward': [1,0,0],\n",
    "    'Rotate Right': [0,0,-0.5],\n",
    "    'Rotate Left': [0,0,0.5],\n",
    "}\n",
    "preds = {}\n",
    "\n",
    "def reset():\n",
    "    x_cond_pixels = x_start\n",
    "    reconstructed_image=x_cond_pixels.to(device)\n",
    "    preds['x_cond_pixels_display'] = (reconstructed_image[-1] * 127.5 + 127.5).clamp(0, 255).permute(1, 2, 0).to(\"cpu\", dtype=torch.uint8).numpy()\n",
    "    preds['x_cond_pixels'] = x_cond_pixels\n",
    "    preds['video'] = [preds['x_cond_pixels_display']]\n",
    "    \n",
    "\n",
    "reset()\n",
    "display(Image.fromarray(preds['x_cond_pixels_display']))\n",
    "Image.fromarray(preds['x_cond_pixels_display']).save('sacson.png')\n",
    "\n",
    "output = widgets.Output()\n",
    "display(output)\n",
    "rel_t = (torch.ones(1)*0.0078125).to(device)\n",
    "\n",
    "@output.capture()\n",
    "def update_image(b):\n",
    "    \n",
    "    if b.description == 'Reset':\n",
    "        print(\"Reset clicked!\")\n",
    "        output.clear_output(wait=False)\n",
    "        reset()\n",
    "        return\n",
    "\n",
    "    print(\"Button clicked!\")\n",
    "    y = commands[b.description]\n",
    "    y = torch.tensor(y).to(device).unsqueeze(0)\n",
    "    \n",
    "    print(\"You entered:\", b.description)\n",
    "    x_cond_pixels = preds['x_cond_pixels'][-4:].unsqueeze(0).to(device)\n",
    "    samples = model_forward_wrapper((model, diffusion, vae), x_cond_pixels, y, None, latent_size, device, config[\"context_size\"], num_goals=1, rel_t=rel_t, progress=True)\n",
    "    x_cond_pixels = samples # torch.clip(samples, -1., 1.)\n",
    "    preds['x_cond_pixels'] = torch.cat([preds['x_cond_pixels'].to(x_cond_pixels), x_cond_pixels], dim=0)\n",
    "    samples = (samples * 127.5 + 127.5).permute(0, 2, 3, 1).clamp(0,255).to(\"cpu\", dtype=torch.uint8).numpy()\n",
    "    display(Image.fromarray(samples[0]))\n",
    "    preds['video'].append(samples[0])\n",
    "\n",
    "buttons = []\n",
    "for o in [\"Forward\", \"Rotate Left\", \"Rotate Right\", \"Reset\"]:\n",
    "    b = widgets.Button(description=o)\n",
    "    b.on_click(update_image)\n",
    "    display(b)\n",
    "    buttons.append(b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c695d6bf-7845-415f-b430-a26e7f37a439",
   "metadata": {},
   "source": [
    "# Generate a video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73a55d72-1a1d-4122-a52f-62e4a4ed6f9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import animation\n",
    "from IPython.display import HTML\n",
    "import numpy as np\n",
    "# np array with shape (frames, height, width, channels)\n",
    "video = np.array(preds['video']) \n",
    "\n",
    "fig = plt.figure()\n",
    "im = plt.imshow(video[0,:,:,:])\n",
    "\n",
    "plt.close() # this is required to not display the generated image\n",
    "\n",
    "def init():\n",
    "    im.set_data(video[0,:,:,:])\n",
    "\n",
    "def animate(i):\n",
    "    im.set_data(video[i,:,:,:])\n",
    "    return im\n",
    "\n",
    "anim = animation.FuncAnimation(fig, animate, init_func=init, frames=video.shape[0],\n",
    "                               interval=500)\n",
    "HTML(anim.to_html5_video())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fd48cd5-84d0-455a-91cd-dd38f97a9ec4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# optional: load from dataset\n",
    "# dataloaders = {}\n",
    "\n",
    "# for dataset_name in config[\"datasets\"]:\n",
    "#     data_config = config[\"datasets\"][dataset_name]\n",
    "#     for data_split_type in [\"test\"]: #[\"test\"]:\n",
    "#         dataset = TrainingDataset(\n",
    "#             data_folder=data_config[\"data_folder\"],\n",
    "#             data_split_folder=data_config[data_split_type],\n",
    "#             dataset_name=dataset_name,\n",
    "#             image_size=config[\"image_size\"],\n",
    "#             min_dist_cat=config[\"distance\"][\"min_dist_cat\"],\n",
    "#             max_dist_cat=config[\"distance\"][\"max_dist_cat\"],\n",
    "#             len_traj_pred=config[\"len_traj_pred\"],\n",
    "#             context_size=config[\"context_size\"],\n",
    "#             normalize=config[\"normalize\"],\n",
    "#             goals_per_obs=1,\n",
    "#             transform=transform,\n",
    "#             predefined_index=None,\n",
    "#             traj_stride=1,\n",
    "#         )\n",
    "#         dataloaders[f\"{dataset_name}_{data_split_type}\"] = dataset\n",
    "#         print(f\"Dataset: {dataset_name} ({data_split_type}), size: {len(dataset)}\")\n",
    "\n",
    "# load from dataset\n",
    "# ds = dataloaders['recon_test'] # scand_test, \n",
    "# x, _, _ = ds[np.random.randint(len(ds))]\n",
    "# x_start = x[:config[\"context_size\"]]\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (nwm)",
   "language": "python",
   "name": "myenv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
